{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Breaking Software ECC with TraceWhisperer *and* SAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "The [uecc_part1_trace.ipynb](uecc_part1_trace.ipynb) notebook was written to show how to find and exploit side-channel leakage in the [micro-ecc library](https://github.com/newaetech/chipwhisperer/tree/develop/hardware/victims/firmware/crypto/micro-ecc) library using Arm trace and [TraceWhisperer](https://github.com/newaetech/tracewhisperer).\n",
    "\n",
    "Most of what makes that notebook long is dealing with the jitter that's inherent to trace.\n",
    "\n",
    "While it's not explicitely presented as such, the notebook essentially deals with the jitter by using SAD (Sum of Absolute Differences) to re-align the traces. There, the SADs are computed in software. But Husky can do SAD in hardware (and it's *much* faster than software-based SAD).\n",
    "\n",
    "At the time that uecc_part1_trace.ipynb was written, it wasn't possible to combine trace and SAD. But with the addition of sequenced triggers, it is now.\n",
    "\n",
    "The [uecc_part2_notrace.ipynb](uecc_part2_notrace.ipynb) notebook already showed a successful attack using SAD. So why this notebook? The answer is that combining trace and SAD **simplifies** the attack:\n",
    "\n",
    "1. **Combining trace and SAD allows us to deal with jitter much more easily than what we had to do in part 1.**\n",
    "2. **Combining trace and SAD makes SAD much easier to tune than it was in part 2.**\n",
    "\n",
    "This notebook and the original uecc_part1_trace.ipynb are similar and independent of one another, but uecc_part1_trace.ipynb has more explanations on the attack and the setting up of trace. It also explores leakage at various points, whereas this notebook focuses on just one area of leakage.\n",
    "\n",
    "Finally, while uecc_part1_trace.ipynb can be run with ChipWhisperer-Pro and PhyWhisperer (running TraceWhisperer FW), this notebook **requires** CW-Husky (for its sequenced trigger capability)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supported Hardware\n",
    "\n",
    "This tutorial requires CW-Husky; it is written for the STM32F3 target, but it could be ported to other Arm targets, as long as they contain an ETM module (which our SAM4S target, sadly, does not)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PLATFORM = 'CW308_STM32F3'\n",
    "TRACE_INTERFACE = 'swo'\n",
    "SCOPETYPE = 'OPENADC'\n",
    "\n",
    "# not supported by this notebook, but can be made to work:\n",
    "#PLATFORM = 'CW308_K82F'\n",
    "#TRACE_INTERFACE = 'parallel'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# platform setup:\n",
    "scope = cw.scope()\n",
    "%run \"../../Setup_Scripts/Setup_Generic.ipynb\"\n",
    "scope.trace.target = target\n",
    "trace = scope.trace\n",
    "scope.clock.clkgen_freq = 10e6\n",
    "scope.clock.clkgen_src = 'system'\n",
    "scope.clock.adc_mul = 1\n",
    "scope.gain.setGain(19)\n",
    "target.baud = 38400 * 10 / 7.37"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.enabled = True\n",
    "trace.clock.clkgen_enabled = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.adc.samples = 6000000\n",
    "scope.adc.stream_mode = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Program STM32 target:\n",
    "\n",
    "**Warning**: if you make any changes to the target firmware (including compiler version and switches), there is a chance that the attack parameters used in this notebook won't work for you anymore. So, for your first run-through, stick with the provided binary.\n",
    "\n",
    "But, making changes to the target firmware is a great way to learn how to use TraceWhisperer, so once you've had success with the default bitfile, do go ahead and try some changes! In fact the TraceWhisperer should make it easier to port the attack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%bash -s \"$PLATFORM\"\n",
    "#cd ../../../firmware/mcu/simpleserial-ecc\n",
    "#make PLATFORM=$1 CRYPTO_TARGET=MICROECC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = '../../../firmware/mcu/simpleserial-ecc/simpleserial-ecc-{}.hex'.format(PLATFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (PLATFORM == 'CW308_STM32F3') or (PLATFORM == 'CWLITEARM'):\n",
    "    prog = cw.programmers.STM32FProgrammer\n",
    "    cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reset_target(scope)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# target info and buildtimes:\n",
    "print(trace.phywhisperer_name())\n",
    "print(trace.get_fw_buildtime())\n",
    "print(scope.fpga_buildtime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set SWO operation mode:\n",
    "\n",
    "Arm processors which support JTAG and SWD come out of reset in JTAG mode. In order to get trace data out of the SWO pin, we need to switch it over to SWD mode.\n",
    "\n",
    "The `jtag_to_swd()` call below runs a special sequence on the TMS and TCK pins to do this switchover. However, different processors may have *additional* requirements to enable the SWO pin. The `simpleserial-trace` firmware handles this for our STM32 target.\n",
    "\n",
    "Another sure-fire way to get a target into SWD mode is to use an external debugger. In that case, do not call `jtag_to_swd()`, as this could result in contention on the TMS/TCK pins, but do call `trace.set_trace_mode()`, because Husky still needs to know that the target is in SWO mode.\n",
    "\n",
    "This table shows the jumper cables that you need to connect between Husky and the target:\n",
    "\n",
    "| ChipWhisperer | Target     |\n",
    "|     :-:       |    :-:     |\n",
    "|      D0       |    TMS     |\n",
    "|      D1       |    TCK     |\n",
    "|      D2       |    TDO     |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRACE_INTERFACE == 'swo':\n",
    "    trace.clock.fe_clock_src = 'target_clock'\n",
    "    assert trace.clock.fe_clock_alive, \"Hmm, the clock you chose doesn't seem to be active.\"\n",
    "    trace.trace_mode = 'SWO'\n",
    "    trace.jtag_to_swd() # switch target into SWO mode\n",
    "\n",
    "    # Now the complicated bit:\n",
    "    acpr = 0\n",
    "    trigger_freq_mul = 8\n",
    "    trace.clock.swo_clock_freq = scope.clock.clkgen_freq * trigger_freq_mul\n",
    "    trace.target_registers.TPI_ACPR = acpr\n",
    "    trace.swo_div = trigger_freq_mul * (acpr + 1)\n",
    "    assert trace.clock.swo_clock_locked, \"Trigger/UART clock not locked\"\n",
    "    assert scope.userio.status & 0x4, \"SWO line not high\"\n",
    "\n",
    "else:\n",
    "    print(\"Not supported in this notebook. See TraceWhisperer.ipynb to see how to set this up.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.clock.reset_adc()\n",
    "time.sleep(0.2)\n",
    "assert (scope.clock.adc_locked), \"ADC failed to lock\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check that the target is alive:\n",
    "If `get_fw_buildtime()` produces no output, the target may have become unresponsive after the above changes; it may simply require a reset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reset_target(scope)\n",
    "print(trace.get_fw_buildtime())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trigger trace capture from target FW:\n",
    "(refer to [uecc_part1_trace.ipynb](uecc_part1_trace.ipynb) for explanations on what this does)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.module = 'basic'\n",
    "scope.trigger.triggers = 'tio4'\n",
    "trace.capture.trigger_source = 'firmware trigger'\n",
    "trace.capture.raw = False\n",
    "\n",
    "# match on any PC match (isync) trace packet:\n",
    "trace.set_pattern_match(0, [3, 8, 32, 0, 0, 0, 0, 0], [255, 255, 255, 0, 0, 0, 0, 0])\n",
    "\n",
    "# enable matching rule:\n",
    "trace.capture.rules_enabled = [0]\n",
    "\n",
    "trace.capture.mode = 'while_trig'\n",
    "\n",
    "TRACES = 'HARDWARE'\n",
    "%run \"ECC_capture.ipynb\"\n",
    "\n",
    "trace.target_registers.DWT_CTRL = '40000021'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing we do different from uecc_part1_trace.ipynb is that we'll only look at one PC address match: the start of the `XYcZadd()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.set_isync_matches(addr0=0, addr1=0x080011bc, match=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "def new_point():\n",
    "    tries = 100\n",
    "    for i in range(tries):\n",
    "        x = random.getrandbits(256)\n",
    "        y = curve.y_recover(x)\n",
    "        if y:\n",
    "            return (x,y)\n",
    "    raise ValueError('Failed to generate a random point')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin the attack by collecting a single trace of the full target operation. We'll use a $k$ with alternating ones/zeros to make things easier for us later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kr = 0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n",
    "k = input_k(kr)\n",
    "Px, Py = new_point()\n",
    "\n",
    "trace.arm_trace()\n",
    "ptrace = capture_ecc_trace(k, Px, Py)\n",
    "while trace.fifo_empty(): pass\n",
    "raw = trace.read_capture_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we get the trace match event timestamps; there should be 255, because that's the number of iterations in the target's main loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times_p2 = trace.get_rule_match_times(raw, rawtimes=False, verbose=False)\n",
    "assert len(times_p2) == 255"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's overlay plot all the power trace segments using the debug trace timestamps.\n",
    "\n",
    "If we're lucky, we may see nice alignment for part of the capture, but this depends on how jittery the target's trace module is feeling at the moment(!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = -200\n",
    "samples = 200\n",
    "from bokeh.palettes import inferno\n",
    "from bokeh.plotting import figure, show\n",
    "from bokeh.resources import INLINE\n",
    "from bokeh.io import output_notebook\n",
    "from bokeh.models import Span, Legend, LegendItem\n",
    "import itertools\n",
    "output_notebook(INLINE)\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(255))\n",
    "for i in range(255):\n",
    "    B.line(list(range(samples)), ptrace.wave[times_p2[i][0]+start:times_p2[i][0]+start+samples], color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In case of bad jitter, here's just one trace segment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = 200\n",
    "start = -200\n",
    "s = figure(width=2000)\n",
    "i1 = 127\n",
    "s.line(list(range(samples)), ptrace.wave[times_p2[i1][0]+start:times_p2[i1][0]+start+samples], line_color='blue')\n",
    "#s.line(range(samples), ptrace.wave[times_p2[i2][0]+start:times_p2[i2][0]+start+samples], line_color='red')\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the trace jitter doesn't act up, what we find is alignment across the power trace segments from sample ~40 to ~170; YMMV due to jitter, but what you're looking for is a series of 7 equidistant narrow peaks:\n",
    "![7_narrow_peaks](img/uecc_7narrow_peaks.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you spend some time studying the power trace, you'll see that this is a particularly nice distinct pattern in the power trace that's probably a pretty good candidate for a SAD reference (spoiler: it is).\n",
    "\n",
    "But there's one small problem: notice how we have indexed into the full power trace 200 samples *before* the trace trigger (i.e. see the `start = -200` above)?\n",
    "\n",
    "Remember our goal is to sequence two triggers: first the trace trigger brings us in the vicinity of the leakage we wish to exploit (+/- some jitter), then a SAD trigger shortly after the trace trigger eliminates the jitter. The key word here is **after**: we can't use the 7 narrow peaks for SAD if they come *before* the trace trigger!\n",
    "\n",
    "What to do? Well, with trace with can trigger on *any* PC address; what if we trace trigger a little bit earlier?\n",
    "\n",
    "Recall that here we were triggering on the start of the `XYcZ_add()` function. Examine the source disassembly, let's move the PC trigger to the previous function call, which is the call to `uECC_vli_set()` that's done at the end of `XYcZ_addC()`. Let's see what happens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.set_isync_matches(addr0=0, addr1=0x08000e12, match=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.arm_trace()\n",
    "ptrace = capture_ecc_trace(k, Px, Py)\n",
    "while trace.fifo_empty(): pass\n",
    "raw = trace.read_capture_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert debug traces into timestamps:\n",
    "times_p2 = trace.get_rule_match_times(raw, rawtimes=False, verbose=False)\n",
    "assert len(times_p2) == 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get 256 PC match events -- that's good.\n",
    "\n",
    "Overlaying all the trace segments, there seems to be more jitter, and our 7 narrow peaks are harder to locate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = 0\n",
    "samples = 1000\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(255))\n",
    "for i in range(255):\n",
    "    B.line(list(range(samples)), ptrace.wave[times_p2[i][0]+start:times_p2[i][0]+start+samples], color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if we plot just a few (or, in the worst case, just one) trace segments, then we should easily locate the 7 peaks near index 500:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = 1000\n",
    "start = 0\n",
    "s = figure(width=2000)\n",
    "# unless you're lucky you'll likely need to try different indices until you get synchronized traces:\n",
    "i1 = 127\n",
    "i2 = 5\n",
    "i3 = 200\n",
    "s.line(list(range(samples)), ptrace.wave[times_p2[i1][0]+start:times_p2[i1][0]+start+samples], line_color='blue')\n",
    "#s.line(list(range(samples)), ptrace.wave[times_p2[i2][0]+start:times_p2[i2][0]+start+samples], line_color='red')\n",
    "s.line(list(range(samples)), ptrace.wave[times_p2[i3][0]+start:times_p2[i3][0]+start+samples], line_color='green')\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's establish our SAD reference. Zoom in to gather good indices.\n",
    "\n",
    "We'll use a 96-sample SAD reference.\n",
    "\n",
    "*(Side note: why 96? It's historical... When this notebook was originally developed, we were limited to setting the SAD reference to 192 or 96 samples. Now Husky SAD is much more flexible and we can set the refence to any number of samples (up to scope.SAD.sad_reference_length), but we'll stay with the original 96 samples.)*\n",
    "\n",
    "\n",
    "Because of the trace jitter, you may need to adjust these numbers. Typically, either 510 or 570 is a good stop point. **Make sure to not go more than 50 samples beyond the end of the 7 peaks.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segment = i1\n",
    "#stop = 510\n",
    "stop = 570\n",
    "start = stop-96\n",
    "\n",
    "ref_trace = ptrace.wave[times_p2[segment][0]+start:times_p2[segment][0]+start+scope.SAD.sad_reference_length*2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = figure(width=2000)\n",
    "s.line(list(range(96)), ref_trace[:96])\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If you don't see the 7 peaks, try the \"other\" value for `stop` (510 or 570).**\n",
    "\n",
    "Recall that our eventual goal is to do a sequenced trace+SAD capture, but let's build that one step at a time.\n",
    "\n",
    "First, let's do a SAD-triggered capture only (without trace). This is helpful for establishing a good SAD threshold value.\n",
    "\n",
    "Instead of capturing the full target operation, we'll use Husky's segmented capture feature, to capture `scope.adc.samples` every time the SAD trigger fires (which we expect it to do 256 times, ideally)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.module = 'SAD'\n",
    "scope.SAD.trigger_sample = 96\n",
    "scope.SAD.reference = ref_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to explicitely tell the SAD module that it can fire multiple times; additionally, we set `scope.SAD.always_armed` so that it keeps firing even after the capture is done.\n",
    "\n",
    "Without this, the SAD module would stop firing after the capture is complete (e.g. after it's fired `scope.adc.segments` times). This way, we can find out whether it's firing too few times or too many times (instead of inferring it from the quality of the traces, or whether the attack works or not).\n",
    "\n",
    "We'll also set some threshold values that we'll fine-tune later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.SAD.multiple_triggers = True\n",
    "scope.SAD.always_armed = True\n",
    "scope.SAD.threshold = 10\n",
    "scope.SAD.interval_threshold = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.adc.stream_mode = False\n",
    "scope.adc.segment_cycle_counter_en = False\n",
    "scope.adc.segments = 255"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again we'll use `SADExplorer` to help tune the thresholds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "explorer = cw.SADExplorer(scope, target, ref_trace, 0, max_segments=255, capture_function=lambda: capture_ecc_trace(k, Px, Py))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike in uecc_part2_notrace.ipynb, where we aimed to tune the thresholds to get exactly 255 matches, here we are ok to get many more matches. As long as you get < 4000 matches, you should be good.\n",
    "\n",
    "Make sure that the 7 peaks of our SAD reference are clear. *(hint: you'll likely need to reduce the SAD thresholds; \"trigger too soon\" errors are ok and expected here)*\n",
    "\n",
    "**This is the power of sequencing triggers!**\n",
    "\n",
    "Why? Because the next step is to sequentially trigger from SAD *within a small time window* after the trace trigger.\n",
    "\n",
    "If you don't get between 255 and 4000 triggers, adjust the thresholds until you do; if that doesn't work, make sure your SAD reference is appropriate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 255 < scope.SAD.num_triggers_seen < 4000:\n",
    "    print('Looks good! Got %d triggers. ✅' % scope.SAD.num_triggers_seen)\n",
    "else:\n",
    "    print('❌ Got %d triggers; try again.' % scope.SAD.num_triggers_seen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next: let's do a just trace-triggered segmented capture, without SAD.\n",
    "\n",
    "At the start of this notebook, we used trace, but the scope capture itself was triggered by IO4.\n",
    "\n",
    "Now we set up the trace module to emit the capture trigger. This will be the first trigger of our trigger sequence, so let's make sure we can get it to fire as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scope.trigger.module = 'trace'\n",
    "trace.capture.trigger_source = 0\n",
    "trace.capture.mode = 'count_cycles'\n",
    "trace.capture.count = int(7e6)\n",
    "trace.capture.max_triggers = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trace.arm_trace()\n",
    "ptrace = capture_ecc_trace(k, Px, Py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "while trace.fifo_empty(): pass\n",
    "raw = trace.read_capture_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert debug traces into timestamps:\n",
    "times_p2 = trace.get_rule_match_times(raw, rawtimes=False, verbose=False)\n",
    "assert len(times_p2) == 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "You can plot the segments out of curiosity, but they won't align nicely due to the jitter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "samples = scope.adc.samples\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(255))\n",
    "for i in range(255):\n",
    "    B.line(list(range(samples)), ptrace.wave[i*samples:(i+1)*samples], color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finally, we are ready to set up the trigger sequencer:\n",
    "\n",
    "Now we add the second trigger to our trigger sequence, the SAD trigger. All we need to do is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.sequencer_enabled = True\n",
    "scope.trigger.module[0] = 'trace'\n",
    "scope.trigger.module[1] = 'SAD'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A powerful feature of sequenced triggering is that we can optionally specify the allowed time delta between the first and second triggers.\n",
    "\n",
    "Minimum and maximum deltas can be specified; if the second trigger occurs outside of this time window, then it is ignored.\n",
    "\n",
    "Our reference trace was taken to end 510 samples after the trace trigger (you may have used a different value: if so, adjust the window accordingly).\n",
    "\n",
    "In the case of the SAD trigger, it's important to know that the SAD trigger actually fires `scope.SAD.latency` cycles *after* the end of the SAD pattern (`scope.SAD.latency` can vary as the design and capture hardware evolves).\n",
    "\n",
    "Additionally, if you ran uecc_part1_trace.ipynb then you know that trace jitter is on the order of +/-70 clock cycles.\n",
    "\n",
    "Taking all this into account, we can expect the SAD trigger to fire in the range of `510 + scope.SAD.trigger_sample + scope.SAD.latency +/- 70` cycles.\n",
    "\n",
    "Finally, while the SAD trigger fires at the end of the SAD pattern match, it must be enabled when the pattern started, and so the start of the window should be shifted back by `scope.SAD.sad_reference_length` , plus a bit more margin on either side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "margin = 100\n",
    "scope.trigger.window_start = 510 + scope.SAD.latency - 70 - margin\n",
    "scope.trigger.window_end = 570 + scope.SAD.latency + scope.SAD.trigger_sample + 70 + margin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have a logic analyzer it can be helpful for tuning the trigger sequencer parameters.\n",
    "\n",
    "For this notebook, this should not be necessary at all, since there is lots of guidance for setting the proper triggering parameters. However, it can come in very handy when you're building a trigger sequence from scratch, so it's good to know that this is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.userio.mode = 'swo_trace_plus_debug'\n",
    "scope.userio.fpga_mode = 14"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Husky's USERIO port pulls double duty here: pins D0, D1, and D2 are connected to the target (to obtain the debug trace data); connect the remaining pins (D3-D7) to your logic analyzer.\n",
    "\n",
    "Printing the `scope.userio` object tells you the definition of each USERIO pin:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.userio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `D3/trigger[0]` is the first trigger in the trigger sequence (trace)\n",
    "- `D4/trigger[1]` is the second trigger (SAD)\n",
    "- `D5/trigger 0 window`: when this is high, the trigger sequencer is waiting for the first trigger; it goes low when the trigger is received\n",
    "- `D6/trigger 1 window`: when this is high, the trigger sequencer is waiting for the second trigger; it goes low when the trigger is received, or its expected window expires\n",
    "- `D7/too late` pulses if the second trigger is not received by the end of its window\n",
    "\n",
    "***Important note**: if you connect a logic analyzer to the USERIO D3-D7 pins, be sure to connect several ground lines between the logic analyzer and Husky. The internal trigger signals are narrow single-cycle pulses, and less-than-ideal connections can actually mess up their proper functioning inside the FPGA (i.e. this can cause the capture that follows to fail). If you suspect this is a problem (i.e. you can't get the sequence-triggered capture to work), try unconnecting D3-D7, and set `scope.userio.mode` back to `'trace'`.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to capture. Some tuning of `scope.SAD.threshold` may still be necessary here; if anything, you **may** find that you need to set the threshold a fair bit higher.\n",
    "\n",
    "This is ok because `scope.trigger.window_start` and `scope.trigger.window_end` prevent the SAD trigger from firing when it's not supposed to."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we no longer need to capture trace data; we're only using trace to generate the first trigger in our sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.capture.mode = 'off'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also turn this off, since SAD will only be allowed to fire in our specified window anyhow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.SAD.always_armed = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The moment of truth: let's see if our sequenced trigger capture works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cross your fingers...\n",
    "trace.arm_trace()\n",
    "seqtrace = capture_ecc_trace(k, Px, Py)\n",
    "assert scope.SAD.num_triggers_seen == 255, 'Got %d SAD triggers' % scope.SAD.num_triggers_seen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert debug traces into timestamps:\n",
    "times_p2 = trace.get_rule_match_times(raw, rawtimes=False, verbose=False)\n",
    "assert len(times_p2) == 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It shouldn't be possible to get too many triggers (because now the SAD can only trigger after the trace trigger), but you may get too few; if so, tweak `scope.SAD.threshold` and/or `scope.SAD.interval_threshold` until you get the right number.\n",
    "\n",
    "Once you do get the right number, let's check whether the time deltas between successive triggers is within the accepted range:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ttimes = scope.trigger.get_trigger_times()\n",
    "assert len(ttimes) == 254\n",
    "assert 20000 < min(ttimes) < 23000\n",
    "assert 20000 < max(ttimes) < 23000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(ttimes), min(ttimes), max(ttimes), np.average(ttimes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're getting out-of-range trigger times, it's likely that you've swung too far on `scope.SAD.threshold` and/or `scope.SAD.interval_threshold`; reduce until you get 255 SAD triggers that are all in the expected time range."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we now overlay the power trace segments, we should find perfect alignment. Compare this with the jittery trace-triggered power trace segments!\n",
    "\n",
    "**If you don't have perfect alignment, it's likely that your SAD threshold is too high: it's very important to fix this before proceeding any further.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = scope.adc.samples\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(255))\n",
    "for i in range(255):\n",
    "    B.line(list(range(samples)), seqtrace.wave[i*samples:(i+1)*samples], color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All the hard work is done now!\n",
    "\n",
    "We have rock-steady, jitter-free captures around where we expect to find leakage, so all that's left to do is find the leakage and exploit it.\n",
    "\n",
    "As in uecc_part1_trace.ipynb, we'll capture a few traces using a constant $k$, calculate the average trace segment for $k$ bits that are 0 and for $k$ bits that are one, and hope to find a consistant difference.\n",
    "\n",
    "While we've tuned our SAD parameters for a single capture, power traces are noisy so it's possible that some captures don't work out. But we can detect this and discard bad traces, instead of letting them pollute our trace set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.capture.use_husky_arm = True # this saves us from have to arm the scope *and* the trace module separately"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll increase our `presamples` to ensure we grab the area where leakage is present:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.adc.presamples = scope.SAD.trigger_sample + scope.SAD.latency + 100\n",
    "scope.adc.samples = scope.adc.presamples + 9 + (3 - (scope.adc.presamples%3)) # must be a multiple of 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traces = 30\n",
    "\n",
    "from tqdm.notebook import tnrange\n",
    "\n",
    "ptraces = []\n",
    "\n",
    "# acquire power and debug traces:\n",
    "for t in tnrange(traces, desc='Capturing traces'):\n",
    "    Px, Py = new_point()\n",
    "    #trace.arm_trace() # don't actually need trace data; just its trigger!\n",
    "    ptrace = capture_ecc_trace(k, Px, Py)\n",
    "    # make sure it's a \"good\" trace:\n",
    "    if scope.SAD.num_triggers_seen != 255:\n",
    "        print('Got %d SAD triggers; skipping this one.' % scope.SAD.num_triggers_seen)\n",
    "        continue\n",
    "    ttimes = scope.trigger.get_trigger_times()\n",
    "    if not ((20000 < min(ttimes) < 23000) and (20000 < max(ttimes) < 23000)):\n",
    "        print('ttimes out of spec: min=%d, max=%d; skipping this one.' % (min(ttimes), max(ttimes)))\n",
    "        continue\n",
    "    ptraces.append(ptrace)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few failed captures is ok, but if you get a lot it's best to fix that by tweaking the SAD thresholds before proceeding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(ptraces) > 25, 'got only %d traces ' % len(ptraces)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute the average ones and zeros:\n",
    "\n",
    "We now take the same approach from part 1 and part 2 to find the leakage and carry out the attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = scope.adc.samples\n",
    "\n",
    "avg_trace = np.zeros(samples)\n",
    "\n",
    "for t in ptraces:\n",
    "    for i in range(1,255):\n",
    "        avg_trace += t.wave[i*samples:(i+1)*samples]\n",
    "\n",
    "avg_trace /= (255*len(ptraces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_ones = np.zeros(samples)\n",
    "avg_zeros = np.zeros(samples)\n",
    "\n",
    "for t in ptraces:\n",
    "    for i in range(254):\n",
    "        if i%2:\n",
    "            avg_ones += t.wave[i*samples:(i+1)*samples]\n",
    "        else:\n",
    "            avg_zeros += t.wave[i*samples:(i+1)*samples]\n",
    "\n",
    "avg_ones /= (127*len(ptraces))\n",
    "avg_zeros /= (127*len(ptraces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = figure(width=2000)\n",
    "\n",
    "xrange = list(range(len(avg_trace)))\n",
    "#s.line(xrange, avg_trace-50, line_color=\"black\")\n",
    "s.line(xrange, avg_ones-50, line_color=\"red\")\n",
    "s.line(xrange, avg_zeros-50, line_color=\"blue\")\n",
    "s.line(xrange, (avg_ones - avg_zeros)*100, line_color=\"orange\")\n",
    "\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The figure below shows the leakage that you should find: the tall positive peak and less tall negative peak, about 40 samples (+/- jitter) before the series of 7 narrow peaks.\n",
    "\n",
    "If you don't see this, it may be that you need to increase `scope.adc.presamples`. (This could happen due to the trace jitter.)\n",
    "\n",
    "![poi11](img/uecc_seqtrig_leakage.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in uecc_part1_trace.ipynb, we extract those peaks to form our list of \"points of interest\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if using a different target, adjust as needed:\n",
    "START=0\n",
    "STOP=300\n",
    "PTHRESH = 50/100\n",
    "NTHRESH = 30/100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poi = list(np.where((avg_ones[START:STOP] - avg_zeros[START:STOP]) > PTHRESH)[0] + START)\n",
    "poi.extend(list(-(np.where((avg_ones[START:STOP] - avg_zeros[START:STOP]) < -NTHRESH)[0] + START)))\n",
    "print(poi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert 5 < len(poi) < 10, \"hmm poi doesn't look quite right, adjust your settings and try again\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_sumdata(poi, ptraces, trim=None):\n",
    "    if trim:\n",
    "        samples = trim\n",
    "    else:\n",
    "        samples = scope.adc.samples\n",
    "    sumdata = np.zeros(255)\n",
    "    for i in range(255):\n",
    "        for t in ptraces:\n",
    "            for p in poi:\n",
    "                sample = t.wave[i*samples+abs(p)]\n",
    "                if p >= 0:\n",
    "                    sumdata[i] += sample\n",
    "                else:\n",
    "                    sumdata[i] -= sample\n",
    "    return sumdata/len(ptraces)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we find whether we can recognize $kr$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd = calc_sumdata(poi, ptraces)\n",
    "\n",
    "s = figure(width=2000)\n",
    "\n",
    "xrange = list(range(len(sd)))\n",
    "s.line(xrange, sd, line_color=\"red\", line_width=2)\n",
    "\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should get something that looks like this:\n",
    "\n",
    "![poi11](img/uecc_0xaaaa.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we can **attack!**\n",
    "\n",
    "We use a random $k$ and see whether we can correctly guess this $k$ from the power trace, using our `poi`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = random_k()\n",
    "kr = regularized_k(k)\n",
    "hex(k), hex(kr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traces = 30\n",
    "\n",
    "from tqdm.notebook import tnrange\n",
    "ptraces = []\n",
    "# acquire power and debug traces:\n",
    "for t in tnrange(traces, desc='Capturing traces'):\n",
    "    Px, Py = new_point()\n",
    "    #trace.arm_trace() # don't actually need trace data; just its trigger!\n",
    "    ptrace = capture_ecc_trace(k, Px, Py)\n",
    "    # make sure it's a \"good\" trace:\n",
    "    if scope.SAD.num_triggers_seen != 255:\n",
    "        print('Got %d SAD triggers; skipping this one.' % scope.SAD.num_triggers_seen)\n",
    "        continue\n",
    "    ttimes = scope.trigger.get_trigger_times()\n",
    "    if not ((20000 < min(ttimes) < 23000) and (20000 < max(ttimes) < 23000)):\n",
    "        print('ttimes out of spec: min=%d, max=%d; skipping this one.' % (min(ttimes), max(ttimes)))\n",
    "        continue\n",
    "    ptraces.append(ptrace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(ptraces) > 25, 'Need more traces! (got %d)' % len(ptraces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd = calc_sumdata(poi, ptraces)\n",
    "\n",
    "s = figure(width=2000)\n",
    "\n",
    "xrange = list(range(len(sd)))\n",
    "s.line(xrange, sd, line_color=\"red\", line_width=2)\n",
    "show(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should get a fairly well-defined train of high and low values, without any points that are too close to the middle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attack(poi, straces, trim=None, verbose=True):\n",
    "    sd = calc_sumdata(poi, straces, trim=trim)\n",
    "\n",
    "    # guess all bits from waveform:\n",
    "    guess = ''\n",
    "    for i in range(1,255):\n",
    "        if sd[i] > np.average(sd):\n",
    "            guess += '0'\n",
    "        else:\n",
    "            guess += '1'\n",
    "\n",
    "    # first and last bit are unknown, so enumerate the possibilities:\n",
    "    guesses = []\n",
    "    for first in (['0', '1']):\n",
    "        for last in (['0', '1']):\n",
    "            guesses.append(int(first + guess + last, 2))\n",
    "\n",
    "    kr = regularized_k(k)\n",
    "    wrong_bits = []\n",
    "    if kr in guesses:\n",
    "        if verbose: print('✅ Guessed right!')\n",
    "    else:\n",
    "        for kbit in range(1,254):\n",
    "            if int(guess[kbit-1]) != ((kr >> (255-kbit)) & 1):\n",
    "                wrong_bits.append(255-kbit)\n",
    "        if verbose:\n",
    "            print('Attack failed.')\n",
    "            print('Guesses: %s' % hex(guesses[0]))\n",
    "            print('         %s' % hex(guesses[1]))\n",
    "            print('         %s' % hex(guesses[2]))\n",
    "            print('         %s' % hex(guesses[3]))\n",
    "            print('Correct: %s' % hex(kr))\n",
    "            print('%d wrong bits' % len(wrong_bits))\n",
    "    return wrong_bits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The moment of truth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack(poi, ptraces)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how many traces are needed to correctly guess the key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for attack_traces in range(1, len(ptraces)+1):\n",
    "    print('Attacking with %d traces... ' % attack_traces,  end='')\n",
    "    wrong_bits = attack(poi, ptraces[:attack_traces], None, False)\n",
    "    if wrong_bits:\n",
    "        print('failed, %d wrong bits' % len(wrong_bits))\n",
    "    else:\n",
    "        print('passed ✅')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The attack should succeed with as few as 15 traces; what's even more impressive is that most bits are guessed correctly with just a single trace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (venv39)",
   "language": "python",
   "name": "venv39"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
